Analysis of Workflow of MainDiffusionV1 in verl-omni

Analysis of the main_diffusion_v1 entrance of v1 trainer.

1. Overview of Workflow
2. Stage 1: Configuration Parsing and Process Invocation
- 2.1. Configuration Validation
- 2.2. Overview of Workflow run_diffusion_v1()
3. Stage 2: Trainer Initialization
- 3.1. Default Initialization and Configuration Loading
- 3.2. Initializing Trainer (Trainer Base)
4. Stage 3: Training

The bash examples/......sh file in examples are actually a wrapper over verl_omni.trainer.main_diffusion.py or main_diffusion_v1.py. Hereafter, we take main_diffusion_v1.py as example.

1. Overview of Workflow

The workflow can be approximately decomposed into 4 sequential parts.

Run bash, and start from verl_omni.trainer.main_diffusion_v1
Run run_diffusion_v1() function to parse configurations.
trainer.init() to initialize tokenizers, weights, resource pools, etc.
trainer.fit() to really train the model.

Next, we’ll go through parts 2-4.

2. Stage 1: Configuration Parsing and Process Invocation

2.1. Configuration Validation

Inside main(), we will check 3 conditions.

auto_set_device(config) that initializes the environment, e.g., NPU or GPU
validate_config(config)
validate_attention_consistency(config). This is to ensure that actor’s attn_backend and rollout’s rollout_attn_backend must match.

2.2. Overview of Workflow `run_diffusion_v1()`

Next, we move on to run_diffusion_v1() function. This function sketches the framework of the exact workflow that we first initialize Ray resource manager, initialize trainer, and then fit the model.

# force enable v1
config.transfer_queue.enable = True

# initialize TQ in Ray task runner
tq.init(config.transfer_queue)

# Get which trainer to use from config
trainer_cls = get_diffusion_trainer_cls(config.trainer.v1.trainer_mode)  # say, "sync"
self.trainer = trainer_cls(config=config)

# initialize the trainer
self.trainer.init()

# initialize agent loop manager
self.init_agent_loop_manager()

# invoke the trainer to train the model
self.trainer.fit(self.agent_loop_manager)

Note, the self.trainer here is of any class that inherits verl_omni.trainer.diffusion.v1.trainer_base.PolicyGradientDiffusionTrainerV1 and is registered through a decorator @register_diffusion_trainer("xxxx"), for example, verl_omni.trainer.diffusion.v1.trainer_sync.PolicyGradientDiffusionTrainerV1Sync and verl_omni.trainer.diffusion.v1.trainer_separate_async.PolicyGradientDiffusionTrainerV1SeparateAsync

3. Stage 2: Trainer Initialization

Next, let’s look at how trainers and “agent loop manager” are initialized.

3.1. Default Initialization and Configuration Loading

The __init__() accepts configuration dict and loads key-value pairs as class properties.

3.2. Initializing Trainer (Trainer Base)

The init() method of PolicyGradientDiffusionTrainerV1 class invokes _setup() function. The _setup() does the following things in order.

Initialize tokenizer.

Here
Initialize dataloader.

4. Stage 3: Training

The main logic of training is implemented in fit() method of PolicyGradientDiffusionTrainerV1.